package learning.pcfg.inference;

import java.util.Enumeration;
import java.util.HashMap;

import learning.pcfg.model.Grammar;
import learning.pcfg.model.IProduction;
import learning.pcfg.model.InnerProduction;
import learning.pcfg.model.MNFTransformer;
import learning.pcfg.model.TerminalProduction;

public class IndexedGrammar {

	public enum ProductionType { UNARY, BINARY, TERMINAL, UNKNOWN };
	
	public int START_SYMBOL;
	
	// stores mapping from ids to original non-terminal names, such as "S"
	public TypeSystem nonTerminals;
	
	// we store all productions in a single array
	public Production[] productions;
	
	// the different types of productions span different ranges in this array
	int unaryStart, unaryEnd, binaryStart, binaryEnd, terminalStart, terminalEnd;
	
	public IndexedGrammar(Grammar grammar) {
		
		// transform grammar to mnf
		grammar = MNFTransformer.transform(grammar);
		
		nonTerminals = new TypeSystem();
		
		START_SYMBOL = nonTerminals.get("S");
		
		// count number of unary and binary productions
		int unaryCount = 0, binaryCount = 0;
		for (InnerProduction ip : grammar.innerProductions) {
			if (ip.rhs.length == 1) unaryCount++;
			if (ip.rhs.length == 2) binaryCount++;
		}
		
		int currentUnary = unaryStart = 0;
		int currentBinary = binaryStart = unaryEnd = unaryCount;
		int currentTerminal = terminalStart = binaryEnd = unaryCount + binaryCount;
		productions = new Production[grammar.innerProductions.size() + grammar.terminalProductions.size()];
		terminalEnd = productions.length;
		for (InnerProduction ip : grammar.innerProductions) {
			if (ip.rhs.length == 1)
				productions[currentUnary] =
					new Production(ip, ProductionType.UNARY, currentUnary++, ip.learnable,
							nonTerminals.get(ip.lhs),
							nonTerminals.get(ip.rhs[0]), 0);
			if (ip.rhs.length == 2)
				productions[currentBinary] =
					new Production(ip, ProductionType.BINARY, currentBinary++, ip.learnable,
							nonTerminals.get(ip.lhs),
							nonTerminals.get(ip.rhs[0]), nonTerminals.get(ip.rhs[1]));
		}
		for (TerminalProduction tp : grammar.terminalProductions) {
			productions[currentTerminal] = 
				new Production(tp, ProductionType.TERMINAL, currentTerminal++, true,
						nonTerminals.get(tp.lhs), 0, 0);
		}		
	}
	
	public ProductionType getProductionType(int production) {
		if (unaryStart <= production && production <= unaryEnd)
			return ProductionType.UNARY;
        else if (binaryStart <= production && production < binaryEnd)
            return ProductionType.BINARY;
        else if (terminalStart <= production && production < terminalEnd)
            return ProductionType.TERMINAL;
        else
            return ProductionType.UNKNOWN;
	}
	
	public int[] getRangeByType(ProductionType type) {
		if (type == ProductionType.UNARY) {
			return new int[] { unaryStart, unaryEnd };
		} else if (type == ProductionType.BINARY) {
			return new int[] { binaryStart, binaryEnd };
		} else if (type == ProductionType.TERMINAL) {
			return new int[] { terminalStart, terminalEnd };
		} else return new int[] { 0, 0 };
	}
	
	public int numProductions() {
		return productions.length;
	}
	
	public int numProductions(ProductionType type) {
		int[] range = getRangeByType(type);
		return range[1] - range[0];
	}
	
	public Enumeration<Production> productions() {
		return new Enumeration<Production>() {
			final int l = productions.length;
			int p = 0;
			public boolean hasMoreElements() { return (p < l); }
			public Production nextElement() { return productions[p++]; }
		};
	}

	public Enumeration<Production> productions(ProductionType pt) {
		final int start, end;
		if (pt.equals(ProductionType.BINARY)) {
			start = binaryStart; end = binaryEnd;
		} else if (pt.equals(ProductionType.UNARY)) {
			start = unaryStart; end = unaryEnd;
		} else if (pt.equals(ProductionType.TERMINAL)) {
			start = terminalStart; end = terminalEnd;
		} else { start = 0; end = 0; } 
		return new Enumeration<Production>() {
			int p = start;
			public boolean hasMoreElements() { return (p < end); }
			public Production nextElement() { return productions[p++]; }
		};
	}
	
	public static class TypeSystem {
		HashMap<String, Integer> typeIds = new HashMap<String, Integer>();
		private int nextTypeId = 1;
		
		public int get(String name) {
			Integer iv = typeIds.get(name);
			if (iv == null) {
				iv = nextTypeId++;
				typeIds.put(name, iv);
			}
			return iv;
		}
		
		public int size() {
			return typeIds.size();
		}
	}
	
	public static class Production {
		public IProduction production; // reference to production in non-indexed grammar
		public ProductionType productionType;		
		public int id;
		public boolean learnable;
		public int lhs;
		public int rhs1;
		public int rhs2;
		
		public Production(IProduction production, ProductionType productionType, 
				int id, boolean learnable, int lhs, int rhs1, int rhs2) {
			this.production = production;
			this.productionType = productionType;			
			this.id = id;
			this.learnable = learnable;
			this.lhs = lhs;
			this.rhs1 = rhs1;
			this.rhs2 = rhs2;
		}
	}
}
